import os
import sys
import torch
from diffusers import StableDiffusionPipeline
from modelscope import snapshot_download
import numpy as np
from PIL import Image
from datetime import datetime

def log(message):
    """打印带时间戳的日志"""
    timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    print(f"[{timestamp}] {message}")

def download_model(model_id, cache_dir):
    """下载模型并返回模型路径"""
    log(f"开始下载模型 {model_id}")
    try:
        model_dir = snapshot_download(model_id, cache_dir=cache_dir)
        log(f"模型下载完成，保存在: {model_dir}")
        return model_dir
    except Exception as e:
        log(f"模型下载失败: {str(e)}")
        raise


def main():
    try:
        # 设置路径
        model_base_path = 'E:/workspace/llm/models'
        model_id = 'Skywork/SkyReels-V2-I2V-1.3B-540P'
        # model_id = 'Wan-AI/Wan2.1-T2V-1.3B'

        # 下载模型
        model_dir = download_model(model_id, model_base_path)
        print(model_dir)
        log("处理完成!")
        return 0

    except Exception as e:
        log(f"发生错误: {str(e)}")
        import traceback
        traceback.print_exc()
        return 1

if __name__ == "__main__":
    sys.exit(main())